import pandas as pd
from matplotlib import pyplot as plt

from Cal_MACD import calMACD
import warnings
warnings.filterwarnings("ignore")

stocks_path = r'd:\Python\study_data\stocks\\'


class Position():
    def __init__(self, stock_name, date, cnt):
        self.stock_name = stock_name
        self.buy_price = 0
        self.buy_date = date
        self.sell_price = 0
        self.sell_date = ''
        self.hold_days = 0
        self.profit = 0
        self.profit_rate = 0
        self.hold_cnt = 0
        self.buy_money = 0


# len(date_df) - 300,
def backtest(all_buys, hold_days, hold_max):
    date_df = pd.read_csv('date.csv')['date']
    all_buys_df = pd.read_csv(all_buys, index_col=0, low_memory=False).sort_index()
    position = {}
    initial_cash = 10
    cash = initial_cash
    date_list = []
    cash_list = []
    for i in range(7101, len(date_df)):
        date = date_df.iloc[i]
        for stock in list(position.keys()):
            position[stock].hold_cnt += 1
            if position[stock].hold_cnt == hold_days:
                stock_df = pd.read_csv(stocks_path + stock + '.csv', index_col=0)
                position[stock].sell_price = stock_df.iloc[stock_df.index.get_loc(position[stock].buy_date) + hold_days]['close']
                position[stock].profit_rate = position[stock].sell_price / position[stock].buy_price - 1
                cash = cash + position[stock].buy_money * position[stock].profit_rate
                with open('hold.txt', 'a', encoding='utf-8') as f:
                    f.write(
                        f'{date}, sell {stock}, buy_price {position[stock].buy_price} sell_price {position[stock].sell_price} profit_rate {"%.2f" % (position[stock].profit_rate * 100)}% profit {"%.2f" % (position[stock].buy_money * position[stock].profit_rate)} cash {"%.2f" % cash}\n')
                del position[stock]

        if date in all_buys_df.index:
            stocks = all_buys_df.loc[date].dropna().values
            buy_stocks_num = (len(stocks) + 1) // 2
            buy_stocks = stocks[:buy_stocks_num]
            for stock in buy_stocks:
                if stock not in position.keys():
                    if len(position) < hold_max:
                        stock_df = pd.read_csv(stocks_path + stock + '.csv', index_col=0)
                        if 'bj' in stock:
                            limit = 0.3
                        elif 'sz300' in stock:
                            limit = 0.2
                        elif 'sh688' in stock:
                            limit = 0.2
                        else:
                            limit = 0.1
                        if stock_df.close.pct_change().loc[date] < limit * 0.97 and stock_df.index.get_loc(date) > 200:
                            position[stock] = Position(stock, date, i)
                            position[stock].buy_price = stock_df.loc[date, 'close']
                            position[stock].buy_date = date
                            position[stock].buy_money = cash / hold_max
                            with open('hold.txt', 'a', encoding='utf-8') as f:
                                f.write(f'{date}, buy {stock}, buy_price {position[stock].buy_price}\n')

    return '{:.2f}'.format((pow(cash / 10, 1/5) - 1) * 100)


if __name__ == '__main__':
    # test = 'all_buys_DiBeiLi.csv'
    test = 'all_buys_20_day_line.csv'
    # test = 'all_buys_test.csv'
    # test = 'all_buys_macd.csv'
    for j in range(5, 6, 1):
        shou_yi = []
        for i in range(10, 31, 5):
            s = backtest(test, i, j)
            print(f'hold_stocks: {j}, hold_days: {i}, annual_rate: {s}%')
            shou_yi.append(float(s))
        print(f'持股数量: {j}, 平均年化收益率: {sum(shou_yi) / len(shou_yi)}%')
    # s = backtest(test, 20, 5)
    # print(f'annual_rate: {s}%')



